import tushare as ts
import pandas as pd
from typing import Dict, List, Optional
from datetime import datetime, timedelta

class NewsData:
    """消息面数据类
    处理公司公告、行业政策和市场新闻等消息面数据
    """
    
    def __init__(self, ts_token: str):
        """初始化消息面数据类
        
        Args:
            ts_token: Tushare API token
        """
        self.ts_pro = ts.pro_api(ts_token)
    
    def get_company_announcements(self, stock_code: str, start_date: str, end_date: str) -> Optional[Dict]:
        """获取公司公告数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            
        Returns:
            包含公司公告数据的字典
        """
        try:
            df = self.ts_pro.anns(
                ts_code=stock_code,
                start_date=start_date,
                end_date=end_date,
                ann_type='重大事项'
            )
            
            if df.empty:
                return None
                
            # 按公告日期升序排序
            df = df.sort_values('ann_date')
            
            return {
                'ann_date': df['ann_date'].values,
                'title': df['title'].values,
                'content': df['content'].values if 'content' in df.columns else None,
                'type': df['ann_type'].values
            }
        except Exception as e:
            print(f"获取公司公告数据失败: {e}")
            return None
    
    def get_industry_news(self, industry: str, start_date: str, end_date: str) -> Optional[Dict]:
        """获取行业新闻数据
        
        Args:
            industry: 行业名称
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            
        Returns:
            包含行业新闻数据的字典
        """
        try:
            df = self.ts_pro.news(
                src='新浪财经',
                start_date=start_date,
                end_date=end_date
            )
            
            if df.empty:
                return None
            
            # 筛选相关行业新闻
            industry_news = df[df['content'].str.contains(industry, na=False)]
            if industry_news.empty:
                return None
                
            # 按发布时间升序排序
            industry_news = industry_news.sort_values('datetime')
            
            return {
                'datetime': industry_news['datetime'].values,
                'title': industry_news['title'].values,
                'content': industry_news['content'].values,
                'source': industry_news['src'].values
            }
        except Exception as e:
            print(f"获取行业新闻数据失败: {e}")
            return None
    
    def get_market_news(self, start_date: str, end_date: str, keywords: List[str] = None) -> Optional[Dict]:
        """获取市场新闻数据
        
        Args:
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            keywords: 关键词列表，用于筛选新闻
            
        Returns:
            包含市场新闻数据的字典
        """
        try:
            df = self.ts_pro.news(
                src='新浪财经',
                start_date=start_date,
                end_date=end_date
            )
            
            if df.empty:
                return None
                
            # 如果提供了关键词，进行筛选
            if keywords:
                mask = df['content'].str.contains('|'.join(keywords), na=False)
                df = df[mask]
                
            if df.empty:
                return None
                
            # 按发布时间升序排序
            df = df.sort_values('datetime')
            
            return {
                'datetime': df['datetime'].values,
                'title': df['title'].values,
                'content': df['content'].values,
                'source': df['src'].values
            }
        except Exception as e:
            print(f"获取市场新闻数据失败: {e}")
            return None
    
    def get_policy_news(self, start_date: str, end_date: str) -> Optional[Dict]:
        """获取政策相关新闻数据
        
        Args:
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            
        Returns:
            包含政策新闻数据的字典
        """
        try:
            # 政策相关关键词
            policy_keywords = ['政策', '规定', '条例', '办法', '通知', '规划', '意见', '决定']
            
            return self.get_market_news(
                start_date=start_date,
                end_date=end_date,
                keywords=policy_keywords
            )
        except Exception as e:
            print(f"获取政策新闻数据失败: {e}")
            return None